package org.shj.spark.streaming;

import java.io.File;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function0;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.api.java.function.VoidFunction2;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.Time;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaPairDStream;
import org.apache.spark.streaming.api.java.JavaReceiverInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.util.LongAccumulator;

import com.google.common.io.Files;

import scala.Tuple2;


public class JavaRecoverableNetworkWordCount {

	public static void main(String[] args) throws Exception{
		if (args.length != 4) {
	      System.err.println("You arguments were " + Arrays.asList(args));
	      System.err.println(
	          "Usage: JavaRecoverableNetworkWordCount <hostname> <port> <checkpoint-directory>\n" +
	          "     <output-file>. <hostname> and <port> describe the TCP server that Spark\n" +
	          "     Streaming would connect to receive data. <checkpoint-directory> directory to\n" +
	          "     HDFS-compatible file system which checkpoint data <output-file> file to which\n" +
	          "     the word counts will be appended\n" +
	          "\n" +
	          "In local mode, <master> should be 'local[n]' with n > 1\n" +
	          "Both <checkpoint-directory> and <output-file> must be absolute paths");
	      System.exit(1);
	    }
		
		final String ip = args[0];
	    final int port = Integer.parseInt(args[1]);
	    final String checkpointDir = args[2];
	    final String outputPath = args[3];
	    final File outputFile = new File(outputPath);
		
		Function0<JavaStreamingContext> factory = new Function0<JavaStreamingContext>() {
			private static final long serialVersionUID = -4370695619014784773L;

			@Override
			public JavaStreamingContext call() throws Exception {
				return createContext(ip, port, checkpointDir, outputPath);
			}
		};
		
		//可以从失败中恢复Driver，不过还需要指定Driver这个进程运行在cluster模式中，并且在提交应用程序的时候
		//指定supervise
		JavaStreamingContext jsc = JavaStreamingContext.getOrCreate(checkpointDir, factory);
		
		JavaReceiverInputDStream<String> lines = jsc.socketTextStream(ip, port);
		
		JavaDStream<String> words = lines.flatMap(new FlatMapFunction<String, String>() {
			private static final long serialVersionUID = -1824228804906565097L;

			public Iterator<String> call(String line) throws Exception {
				return Arrays.asList(line.split("\\s+")).iterator();
			}
		});
		
		JavaPairDStream<String, Integer> wordCnt = words.mapToPair(new PairFunction<String, String, Integer>() {
			private static final long serialVersionUID = 5998492885266269233L;

			public Tuple2<String, Integer> call(String t) throws Exception {
				return new Tuple2<String,Integer>(t, 1);
			}
		}).reduceByKey(new Function2<Integer, Integer, Integer>() {
			private static final long serialVersionUID = -1674734625587228837L;

			@Override
			public Integer call(Integer v1, Integer v2) throws Exception {
				return v1 + v2;
			}
		});
		
		wordCnt.foreachRDD(new VoidFunction2<JavaPairRDD<String, Integer>, Time>() {
			private static final long serialVersionUID = -465897857226871905L;

			public void call(JavaPairRDD<String, Integer> rdd, Time time) throws Exception {
				// Get or register the blacklist Broadcast
		        final Broadcast<List<String>> blacklist =
		            JavaWordBlacklist.getInstance(new JavaSparkContext(rdd.context()));
		        
		        // Get or register the droppedWordsCounter Accumulator
		        final LongAccumulator droppedWordsCounter =
		            JavaDroppedWordsCounter.getInstance(new JavaSparkContext(rdd.context()));
		        
		        // Use blacklist to drop words and use droppedWordsCounter to count them
		        String counts = rdd.filter(new Function<Tuple2<String, Integer>, Boolean>() {
				    private static final long serialVersionUID = 1L;

				    @Override
		            public Boolean call(Tuple2<String, Integer> wordCount) {
		                if (blacklist.value().contains(wordCount._1)) {
		                  droppedWordsCounter.add(wordCount._2);
		                  return false;
		                } else {
		                  return true;
		                }
		            }
		        }).collect().toString();
		        String output = "Counts at time " + time + " " + counts;
		        System.out.println(output);
		        System.out.println("Dropped " + droppedWordsCounter.value() + " word(s) totally");
		        System.out.println("Appending to " + outputFile.getAbsolutePath());
		        Files.append(output + "\n", outputFile, Charset.defaultCharset());
				
			}
		});
		
		jsc.start();
		jsc.awaitTermination();
		jsc.stop();
	}
	
	private static JavaStreamingContext createContext(String hostname, int port, 
			String checkpointDir, String outputPath) {
		System.out.println("Create new JavaStreamingContext");
		final File outputFile = new File(outputPath);
	    if (outputFile.exists()) {
	      outputFile.delete();
	    }
	    SparkConf conf = new SparkConf().setAppName("BlacklistWordCount");
	    JavaStreamingContext jsc = new JavaStreamingContext(conf, Durations.seconds(10));
	    jsc.checkpoint(checkpointDir);
	    return jsc;
	}
	

}

class JavaWordBlacklist {

  private static volatile Broadcast<List<String>> instance = null;

  public static Broadcast<List<String>> getInstance(JavaSparkContext jsc) {
    if (instance == null) {
      synchronized (JavaWordBlacklist.class) {
        if (instance == null) {
          List<String> wordBlacklist = Arrays.asList("a", "b", "c");
          instance = jsc.broadcast(wordBlacklist);
        }
      }
    }
    return instance;
  }
}

/**
 * Use this singleton to get or register an Accumulator.
 */
class JavaDroppedWordsCounter {

  private static volatile LongAccumulator instance = null;

  public static LongAccumulator getInstance(JavaSparkContext jsc) {
    if (instance == null) {
      synchronized (JavaDroppedWordsCounter.class) {
        if (instance == null) {
          instance = jsc.sc().longAccumulator("WordsInBlacklistCounter");
        }
      }
    }
    return instance;
  }
}


